import pygame
import heapq

# 初始化 Pygame
pygame.init()

# 设置窗口和网格参数
WIDTH, HEIGHT = 600, 850  # 增加高度以容纳第三个输入框
GRID_SIZE = 10
CELL_SIZE = WIDTH // GRID_SIZE
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Manhattan Path with Fixed Directions and Dynamic Turns")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
GRAY = (200, 200, 200)

# 曼哈顿距离计算函数
def manhattan_distance(p1, p2):
    return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

# 计算路径中的转折次数
def count_turns(path):
    if len(path) < 3:
        return 0
    turns = 0
    for i in range(1, len(path) - 1):
        prev = (path[i-1][0] - path[i][0], path[i-1][1] - path[i][1])
        next = (path[i][0] - path[i+1][0], path[i][1] - path[i+1][1])
        if prev != next:
            turns += 1
    return turns

# A* 算法，严格固定起点和终点方向
def a_star(start, goal, target_turns=3, start_dir=(0, 1), end_dir=(0, -1)):
    open_set = [(0, 0, start, [start])]  # (f_score, turns, position, path)
    visited = set()

    while open_set:
        f_score, turns, current, path = heapq.heappop(open_set)
        
        if current == goal and turns == target_turns:
            if len(path) >= 2 and (path[-2][0] - current[0], path[-2][1] - current[1]) == end_dir:
                return path
            continue

        if (current, turns) in visited:
            continue
        visited.add((current, turns))
        
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
        if len(path) == 1:
            directions = [start_dir]
        
        for dx, dy in directions:
            neighbor = (current[0] + dx, current[1] + dy)
            
            if (0 <= neighbor[0] < GRID_SIZE and 
                0 <= neighbor[1] < GRID_SIZE and 
                neighbor not in path):
                
                new_path = path + [neighbor]
                new_turns = count_turns(new_path)
                
                if new_turns > target_turns:
                    continue
                
                if (neighbor, new_turns) not in visited:
                    g_score = len(new_path) - 1
                    h_score = manhattan_distance(neighbor, goal)
                    f_score = g_score + h_score
                    heapq.heappush(open_set, (f_score, new_turns, neighbor, new_path))
    
    return []

# 将字符串转换为方向元组
def str_to_dir(text):
    text = text.lower().strip()
    if text == "right":
        return (0, 1)
    elif text == "down":
        return (1, 0)
    elif text == "left":
        return (0, -1)
    elif text == "up":
        return (-1, 0)
    return None

# 将字符串转换为转折次数
def str_to_turns(text):
    try:
        turns = int(text.strip())
        if turns >= 0:
            return turns
    except ValueError:
        return None
    return None

# 绘制网格和路径
def draw_grid(start, goal, path, start_dir, end_dir, target_turns, start_input, end_input, turns_input, active_input):
    screen.fill(WHITE)
    
    # 绘制网格
    for x in range(0, WIDTH, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (x, 0), (x, WIDTH))
    for y in range(0, WIDTH, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (0, y), (WIDTH, y))
    
    # 绘制路径线段
    if path:
        for i in range(len(path) - 1):
            start_pos = (path[i][1] * CELL_SIZE + CELL_SIZE // 2, path[i][0] * CELL_SIZE + CELL_SIZE // 2)
            end_pos = (path[i+1][1] * CELL_SIZE + CELL_SIZE // 2, path[i+1][0] * CELL_SIZE + CELL_SIZE // 2)
            pygame.draw.line(screen, RED, start_pos, end_pos, 3)
    
    # 如果路径太短，则显示提示
    if len(path) <= 1:  # 如果路径的线段太少
        font = pygame.font.SysFont(None, 30)
        warning_text = font.render("Path too short to draw lines! or target_turns too few ", True, BLACK)
        screen.blit(warning_text, (WIDTH // 2 - warning_text.get_width() // 2, HEIGHT // 2))

    # 绘制起点和终点
    pygame.draw.circle(screen, BLACK, (start[1] * CELL_SIZE + CELL_SIZE // 2, start[0] * CELL_SIZE + CELL_SIZE // 2), CELL_SIZE // 3)
    pygame.draw.circle(screen, GREEN, (goal[1] * CELL_SIZE + CELL_SIZE // 2, goal[0] * CELL_SIZE + CELL_SIZE // 2), CELL_SIZE // 3)
    
    # 绘制输入框
    font = pygame.font.SysFont(None, 30)
    start_label = font.render("Start Dir:", True, BLACK)
    end_label = font.render("End Dir:", True, BLACK)
    turns_label = font.render("Turns:", True, BLACK)
    
    start_rect = pygame.Rect(100, 620, 200, 40)
    end_rect = pygame.Rect(100, 670, 200, 40)
    turns_rect = pygame.Rect(100, 720, 200, 40)
    
    pygame.draw.rect(screen, BLUE if active_input == "start" else GRAY, start_rect, 2)
    pygame.draw.rect(screen, BLUE if active_input == "end" else GRAY, end_rect, 2)
    pygame.draw.rect(screen, BLUE if active_input == "turns" else GRAY, turns_rect, 2)
    
    start_text = font.render(start_input, True, BLACK)
    end_text = font.render(end_input, True, BLACK)
    turns_text = font.render(turns_input, True, BLACK)
    
    screen.blit(start_label, (10, 625))
    screen.blit(end_label, (10, 675))
    screen.blit(turns_label, (10, 725))
    screen.blit(start_text, (start_rect.x + 5, start_rect.y + 5))
    screen.blit(end_text, (end_rect.x + 5, end_rect.y + 5))
    screen.blit(turns_text, (turns_rect.x + 5, turns_rect.y + 5))
    
    # 提示
    hint1 = font.render("Start/End: 'right', 'down', 'left', 'up'", True, BLACK)
    hint2 = font.render("Turns: number (e.g., 1, 2, 3), press Enter", True, BLACK)
    screen.blit(hint1, (10, 770))
    screen.blit(hint2, (10, 800))

# 主程序
def main():
    start = (GRID_SIZE // 2, GRID_SIZE // 2)
    running = True
    target_turns = 3
    start_dir = (0, 1)
    end_dir = (0, -1)
    start_input = "left"
    end_input = "down"
    turns_input = "3"
    active_input = None
    
    clock = pygame.time.Clock()
    
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                mouse_pos = event.pos
                start_rect = pygame.Rect(100, 620, 200, 40)
                end_rect = pygame.Rect(100, 670, 200, 40)
                turns_rect = pygame.Rect(100, 720, 200, 40)
                if start_rect.collidepoint(mouse_pos):  # 修复拼写错误
                    active_input = "start"
                elif end_rect.collidepoint(mouse_pos):
                    active_input = "end"
                elif turns_rect.collidepoint(mouse_pos):
                    active_input = "turns"
                else:
                    active_input = None
            elif event.type == pygame.KEYDOWN and active_input:
                if event.key == pygame.K_RETURN:
                    if active_input == "start":
                        new_dir = str_to_dir(start_input)
                        if new_dir:
                            start_dir = new_dir
                    elif active_input == "end":
                        new_dir = str_to_dir(end_input)
                        if new_dir:
                            end_dir = new_dir
                    elif active_input == "turns":
                        new_turns = str_to_turns(turns_input)
                        if new_turns is not None:
                            target_turns = new_turns
                elif event.key == pygame.K_BACKSPACE:
                    if active_input == "start":
                        start_input = start_input[:-1]
                    elif active_input == "end":
                        end_input = end_input[:-1]
                    elif active_input == "turns":
                        turns_input = turns_input[:-1]
                elif event.unicode.isalnum() or event.unicode in ['r', 'i', 'g', 'h', 't', 'd', 'o', 'w', 'n', 'l', 'e', 'f', 't', 'u', 'p']:
                    if active_input == "start" and len(start_input) < 10:
                        start_input += event.unicode
                    elif active_input == "end" and len(end_input) < 10:
                        end_input += event.unicode
                    elif active_input == "turns" and len(turns_input) < 3 and event.unicode.isdigit():
                        turns_input += event.unicode
        
        # 获取鼠标位置并转换为网格坐标
        mouse_pos = pygame.mouse.get_pos()
        goal = (mouse_pos[1] // CELL_SIZE, mouse_pos[0] // CELL_SIZE)
        goal = (max(0, min(goal[0], GRID_SIZE-1)), max(0, min(goal[1], GRID_SIZE-1)))
        
        # 计算路径
        path = a_star(start, goal, target_turns, start_dir, end_dir)
        
        # 绘制网格和路径
        draw_grid(start, goal, path, start_dir, end_dir, target_turns, start_input, end_input, turns_input, active_input)
        pygame.display.flip()
        
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()
